"""
Some shared tools in evaluation process
"""
import numpy as np
from model.metrics import box_nms, super_nms


# Get the dilation kernel based on the distance threshold
def get_dilation_kernel(dist_thresh):
    # Initialize empty kernel
    kernel_len = round(2 * dist_thresh + 1)
    output_kernel = np.zeros([kernel_len, kernel_len])
    center = np.array([[dist_thresh, dist_thresh]])

    # Intitialize the meshgrid
    xv, yv = np.meshgrid(np.arange(0, kernel_len, 1),
                         np.arange(0, kernel_len, 1),
                         indexing='xy')
    coord_map = np.concatenate([yv[..., None], xv[..., None]], axis=-1)
    coord_map_flattened = coord_map.reshape([-1, 2])
    dist = np.sqrt(np.sum((coord_map_flattened - center) ** 2, axis=-1))
    coord_valid = coord_map_flattened[dist <= dist_thresh, :]

    output_kernel[coord_valid[:, 0], coord_valid[:, 1]] = 1

    return output_kernel.astype(np.uint8)


# Convert the exported junctions and heatmaps predictions to uniform format (Either accumulated or single detection).
def convert_raw_exported_predictions(input_data, grid_size=8, detect_thresh=1 / 65, topk=300):
    """
    input_data: the raw data (dict) decoded from the hdf5 dataset
    outputs: dict containing required entries including
        junctions_pred: Nx2 ndarray containing junction predictions after nms.
        heatmap_pred: HxW ndarray containing predicted heatmaps
        valid_mask: HxW ndarray containing the valid mask
    """
    # Check the input_data is from (1) single prediction, or (2) homography adaptation.
    # Homography adaptation raw predictions
    if ("junc_prob_mean" in input_data.keys()) and ("heatmap_prob_mean" in input_data.keys()):
        # Get the junction predictions and convert if to Nx2 format
        junc_prob = input_data["junc_prob_mean"]
        junc_pred_np = junc_prob[None, ...]
        junc_pred_np_nms = super_nms(junc_pred_np, grid_size, detect_thresh, topk)
        junctions = np.where(junc_pred_np_nms.squeeze())
        junc_points_pred = np.concatenate([junctions[0][..., None], junctions[1][..., None]], axis=-1)

        # Get the heatmap predictions
        heatmap_pred = input_data["heatmap_prob_mean"].squeeze()
        valid_mask = np.ones(heatmap_pred.shape, dtype=np.int32)

    # Single predictions
    else:
        # Get the junction point predictions and convert to Nx2 format
        junc_points_pred = np.where(input_data["junc_pred_nms"])
        junc_points_pred = np.concatenate([junc_points_pred[0][..., None], junc_points_pred[1][..., None]], axis=-1)

        # Get the heatmap predictions
        heatmap_pred = input_data["heatmap_pred"]
        valid_mask = input_data["valid_mask"]

    return {
        "junctions_pred": junc_points_pred,
        "heatmap_pred": heatmap_pred,
        "valid_mask": valid_mask
    }


# Convert line score map to line
# Note: different from Nx4 format, here we convert to Nx2x2 format.
def line_score_map_to_segments(junctions, line_score_map):
    line_score_map_tmp = line_score_map.copy()

    output_segments = np.zeros([0, 2, 2])
    output_scores = np.zeros([0, ])
    for idx in range(junctions.shape[0]):
        # if no connectivity, just skip it
        if line_score_map_tmp[idx, :].sum() == 0:
            continue
        # Record the line segment
        else:
            for idx2 in np.where(line_score_map_tmp[idx, :] > 0.)[0]:
                p1 = junctions[idx, :]  # HW format
                p2 = junctions[idx2, :]  # HW format

                single_seg = np.concatenate([p1[None, ...], p2[None, ...]], axis=0)
                output_segments = np.concatenate((output_segments, single_seg[None, ...]), axis=0)

                single_score = line_score_map_tmp[idx, idx2]
                # import ipdb; ipdb.set_trace()
                output_scores = np.concatenate([output_scores, [single_score]])

                # Update line_map
                line_score_map_tmp[idx, idx2] = 0.
                line_score_map_tmp[idx2, idx] = 0.

    return output_segments, output_scores


# Convert line map to line
# Note: different from Nx4 format, here we convert to Nx2x2 format.
def line_map_to_segments(junctions, line_map):
    line_map_tmp = line_map.copy()

    output_segments = np.zeros([0, 2, 2])
    for idx in range(junctions.shape[0]):
        # if no connectivity, just skip it
        if line_map_tmp[idx, :].sum() == 0:
            continue
        # Record the line segment
        else:
            for idx2 in np.where(line_map_tmp[idx, :] == 1)[0]:
                p1 = junctions[idx, :]  # HW format
                p2 = junctions[idx2, :]  # HW format
                single_seg = np.concatenate([p1[None, ...], p2[None, ...]], axis=0)
                output_segments = np.concatenate((output_segments, single_seg[None, ...]), axis=0)

                # Update line_map
                line_map_tmp[idx, idx2] = 0
                line_map_tmp[idx2, idx] = 0

    return output_segments